#!/usr/bin/env python3

import os
import sys
import xml.etree.ElementTree as ET
import csv
import glob
import argparse
from pathlib import Path


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='将SUMO FCD车辆数据文件(XML)转换为CSV格式')
    parser.add_argument('--folder', type=str, help='请输入要处理的文件夹路径：')
    return parser.parse_args()


def convert_fcd_to_csv(xml_file, csv_file):
    """将FCD XML文件转换为CSV格式"""
    print(f"正在处理文件: {xml_file}")
    print(f"输出CSV文件: {csv_file}")
    
    try:
        # 解析XML文件
        tree = ET.parse(xml_file)
        root = tree.getroot()
        
        # 打开CSV文件进行写入
        with open(csv_file, 'w', newline='') as csvfile:
            # 创建CSV写入器
            csv_writer = csv.writer(csvfile)
            
            # 写入标题行
            csv_writer.writerow(['time', 'vehicle_id', 'type', 'x', 'y', 'angle', 'speed', 'pos', 'lane', 'slope'])
            
            # 遍历XML中的每个时间步
            for timestep in root.findall('timestep'):
                time = timestep.get('time')
                
                # 对每个时间步中的每辆车
                for vehicle in timestep.findall('vehicle'):
                    vehicle_id = vehicle.get('id')
                    vehicle_type = vehicle.get('type')
                    x = vehicle.get('x')
                    y = vehicle.get('y')
                    angle = vehicle.get('angle')
                    speed = vehicle.get('speed')
                    pos = vehicle.get('pos')
                    lane = vehicle.get('lane')
                    slope = vehicle.get('slope')
                    
                    # 将车辆数据写入CSV
                    csv_writer.writerow([time, vehicle_id, vehicle_type, x, y, angle, speed, pos, lane, slope])
        
        print(f"转换完成: {csv_file}")
        return True
    
    except Exception as e:
        print(f"转换失败: {e}")
        return False


def main():
    """主函数"""
    args = parse_args()
    
    # 如果命令行参数中提供了文件夹路径，则使用它；否则，提示用户输入
    folder_path = args.folder
    if not folder_path:
        folder_path = input("请输入包含XML文件的文件夹路径: ")
    
    if not folder_path.strip():
        print("错误: 必须提供文件夹路径")
        return 1
    
    # 检查文件夹是否存在
    if not os.path.isdir(folder_path):
        print(f"错误: 文件夹 '{folder_path}' 不存在")
        return 1
    
    # 获取文件夹中的所有XML文件
    xml_files = glob.glob(os.path.join(folder_path, "*.xml"))
    
    if not xml_files:
        print(f"错误: 文件夹 '{folder_path}' 中没有找到XML文件")
        return 1
    
    # 列出所有可用的XML文件供用户选择
    print("\n可用的XML文件:")
    for i, file_path in enumerate(xml_files, 1):
        print(f"{i}. {os.path.basename(file_path)}")
    
    # 让用户选择文件
    choice = input("\n请选择要处理的XML文件 (输入编号或文件名): ")
    
    try:
        # 尝试将输入解析为数字
        choice_num = int(choice)
        if 1 <= choice_num <= len(xml_files):
            xml_file = xml_files[choice_num - 1]
        else:
            print(f"错误: 无效的选择 {choice_num}，应在1到{len(xml_files)}之间")
            return 1
    except ValueError:
        # 如果不是数字，则视为文件名
        xml_file = os.path.join(folder_path, choice)
        if not os.path.exists(xml_file):
            print(f"错误: 文件 '{xml_file}' 不存在")
            return 1
    
    # 创建输出文件名 (与XML文件同名，但扩展名为.csv)
    csv_file = os.path.splitext(xml_file)[0] + '.csv'
    
    # 转换文件
    success = convert_fcd_to_csv(xml_file, csv_file)
    
    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())
